#!/bin/bash
# sshinit.sh - SSH免密登录初始化脚本

# 设置颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 检查参数
if [ $# -lt 1 ]; then
    echo -e "${RED}用法: $0 <用户名@主机地址> [端口]${NC}"
    echo "例如: $0 user@hostname"
    echo "例如: $0 user@hostname 22"
    exit 1
fi

REMOTE_HOST=$1
PORT=${2:-22}

echo -e "${YELLOW}开始配置SSH免密登录到 $REMOTE_HOST${NC}"

# 新增功能：选择网络接口
select_network_interface() {
    echo -e "${YELLOW}检测到以下网络接口:${NC}"
    interfaces=$(ip link show | awk -F': ' '/^[0-9]+: [a-zA-Z]/ {print $2}' | grep -v lo)

    select interface in $interfaces; do
        if [ -n "$interface" ]; then
            SELECTED_INTERFACE=$interface
            echo -e "${GREEN}已选择网络接口: $SELECTED_INTERFACE${NC}"
            break
        else
            echo -e "${RED}无效选择，请重新输入${NC}"
        fi
    done
}

# 新增功能：添加主机路由
add_host_route() {
    local remote_ip=$(getent ahosts ${REMOTE_HOST#*@} | head -n 1 | awk '{print $1}')

    if [ -z "$remote_ip" ]; then
        echo -e "${RED}无法解析主机IP地址${NC}"
        return 1
    fi

    echo -e "${YELLOW}远程主机IP: $remote_ip${NC}"

    # 检查是否有直连路由
    local direct_route=$(ip route show dev $SELECTED_INTERFACE | grep -E "^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+/[0-9]+" | head -n 1)

    if [ -n "$direct_route" ]; then
        echo -e "${YELLOW}检测到接口 $SELECTED_INTERFACE 的网络: $direct_route${NC}"
        read -p "是否通过网关添加路由? (y/n, 默认为n): " use_gateway

        if [[ "$use_gateway" =~ ^[Yy]$ ]]; then
            # 获取选定接口的网关
            gateway=$(ip route show dev $SELECTED_INTERFACE | grep default | awk '{print $3}')

            if [ -z "$gateway" ]; then
                echo -e "${RED}无法获取网关信息${NC}"
                return 1
            fi

            echo -e "${YELLOW}为 $remote_ip 添加路由，通过接口 $SELECTED_INTERFACE 和网关 $gateway${NC}"

            # 添加经过网关的主机路由
            sudo ip route add $remote_ip/32 via $gateway dev $SELECTED_INTERFACE
        else
            echo -e "${YELLOW}为 $remote_ip 添加直连路由，通过接口 $SELECTED_INTERFACE${NC}"

            # 添加直连主机路由
            sudo ip route add $remote_ip/32 dev $SELECTED_INTERFACE
        fi
    else
        # 没有直连路由信息，尝试获取网关
        gateway=$(ip route show dev $SELECTED_INTERFACE | grep default | awk '{print $3}')

        if [ -n "$gateway" ]; then
            echo -e "${YELLOW}检测到网关: $gateway${NC}"
            read -p "是否通过网关添加路由? (y/n, 默认为y): " use_gateway

            if [[ ! "$use_gateway" =~ ^[Nn]$ ]]; then
                echo -e "${YELLOW}为 $remote_ip 添加路由，通过接口 $SELECTED_INTERFACE 和网关 $gateway${NC}"

                # 添加经过网关的主机路由
                sudo ip route add $remote_ip/32 via $gateway dev $SELECTED_INTERFACE
            else
                echo -e "${YELLOW}为 $remote_ip 添加直连路由，通过接口 $SELECTED_INTERFACE${NC}"

                # 添加直连主机路由
                sudo ip route add $remote_ip/32 dev $SELECTED_INTERFACE
            fi
        else
            echo -e "${YELLOW}未检测到网关，添加直连路由${NC}"

            # 添加直连主机路由
            sudo ip route add $remote_ip/32 dev $SELECTED_INTERFACE
        fi
    fi

    if [ $? -eq 0 ]; then
        echo -e "${GREEN}主机路由添加成功${NC}"
    else
        echo -e "${RED}主机路由添加失败${NC}"
        return 1
    fi
}

# 如果系统有多个网络接口，让用户选择一个
interface_count=$(ip link show | awk -F': ' '/^[0-9]+: [a-zA-Z]/ {print $2}' | grep -v lo | wc -l)
if [ $interface_count -gt 1 ]; then
    select_network_interface
    add_host_route
fi

# 1. 检查本地是否存在SSH密钥
if [ ! -f ~/.ssh/id_rsa.pub ]; then
    echo -e "${YELLOW}未找到SSH公钥，正在生成新的SSH密钥对...${NC}"
    ssh-keygen -t rsa -b 2048 -f ~/.ssh/id_rsa -N ""

    if [ $? -ne 0 ]; then
        echo -e "${RED}SSH密钥生成失败${NC}"
        exit 1
    fi

    echo -e "${GREEN}SSH密钥对生成成功${NC}"
else
    echo -e "${GREEN}已存在SSH密钥对${NC}"
fi

# 2. 确保本地SSH目录权限正确
chmod 700 ~/.ssh
chmod 600 ~/.ssh/id_rsa
chmod 644 ~/.ssh/id_rsa.pub

# 3. 将公钥复制到远程主机
echo -e "${YELLOW}正在将公钥复制到远程主机...${NC}"
echo "请在此处输入远程主机的密码："

# 使用ssh-copy-id复制公钥
ssh-copy-id -p $PORT $REMOTE_HOST

if [ $? -ne 0 ]; then
    echo -e "${RED}公钥复制失败，请检查网络连接和账户凭据${NC}"
    exit 1
fi

# 4. 测试免密登录
echo -e "${YELLOW}正在测试免密登录...${NC}"
ssh -p $PORT -o BatchMode=yes -o ConnectTimeout=10 $REMOTE_HOST 'echo "连接成功"'

if [ $? -eq 0 ]; then
    echo -e "${GREEN}免密登录配置成功！${NC}"
    echo -e "${GREEN}现在你可以直接使用 'ssh $REMOTE_HOST' 登录而无需输入密码${NC}"
else
    echo -e "${RED}免密登录测试失败${NC}"
    exit 1
fi

echo -e "${GREEN}SSH免密登录配置完成！${NC}"